from BPRMF.Recommender import Recommender as BPRRecommender
from Evaluator import evaluate
from POP.Recommender import Recommender as POPRecommender
from ICF.Recommender import Recommender as ICFRecommender
from UCF.Recommender import Recommender as UCFRecommender
from typing import Dict, List
from torch.cuda import is_available
from tqdm import tqdm
import numpy as np
import time
import torch
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '1'


class Recommender:
    def __init__(self, hyper_params):
        self.hyper_params = hyper_params

        if 'bpr' in hyper_params['methods']:
            self.bpr_recommender = BPRRecommender(hyper_params)
        if 'pop' in hyper_params['methods']:
            self.pop_recommender = POPRecommender(hyper_params)
        if 'icf' in hyper_params['methods']:
            self.icf_recommender = ICFRecommender(hyper_params)
        if 'ucf' in hyper_params['methods']:
            self.ucf_recommender = UCFRecommender(hyper_params)

        # POP Recommender are used for filling the remain items
        self.bpr_cnt = hyper_params['bpr_cnt']
        self.icf_cnt = hyper_params['icf_cnt']
        self.ucf_cnt = hyper_params['ucf_cnt']
        self.pop_cnt = hyper_params['pop_cnt']

    def recommend(self) -> Dict[int, List[int]]:
        if 'bpr' in self.hyper_params['methods']:
            bpr_result = self.bpr_recommender.recommend(self.bpr_cnt)
        if 'icf' in self.hyper_params['methods']:
            icf_result = self.icf_recommender.recommend(self.icf_cnt)
        if 'ucf' in self.hyper_params['methods']:
            ucf_result = self.ucf_recommender.recommend(self.ucf_cnt)
        if 'pop' in self.hyper_params['methods']:
            pop_result = self.pop_recommender.recommend(self.pop_cnt)

        result: Dict[int, List[int]] = {}

        print('Recommender: Merging results...')
        for user in tqdm(range(self.hyper_params['user_cnt'])):
            cur_result = []

            if 'bpr' in self.hyper_params['methods'] and user in bpr_result:
                cnt = self.hyper_params['methods']['bpr']
                for rec in bpr_result[user]:
                    if rec in cur_result: continue
                    cur_result.append(rec)
                    cnt -= 1

                    if cnt == 0:
                        break
            if 'icf' in self.hyper_params['methods'] and user in icf_result:
                cnt = self.hyper_params['methods']['icf']
                for rec in icf_result[user]:
                    if rec in cur_result: continue
                    cur_result.append(rec)
                    cnt -= 1

                    if cnt == 0:
                        break
            if 'ucf' in self.hyper_params['methods'] and user in ucf_result:
                cnt = self.hyper_params['methods']['ucf']
                for rec in ucf_result[user]:
                    if rec in cur_result: continue
                    cur_result.append(rec)
                    cnt -= 1

                    if cnt == 0:
                        break
            if 'pop' in self.hyper_params['methods']:
                cnt = 500 - len(cur_result)
                for rec in pop_result:
                    if rec in cur_result: continue
                    cur_result.append(rec)
                    cnt -= 1

                    if cnt == 0:
                        break

            result[user] = cur_result
            print(len(cur_result))

        return result


if __name__ == '__main__':
    hyper_params = {
        'datasets': 'datasets/data_small',
        'device': 'cuda' if is_available() else 'cpu',
        'batch_size': 128 if is_available() else 64,
        'epochs': 5 if is_available() else 2,
        'embed_dim': 64,
        'neg_sample_cnt': 300,
        'epsilon': 0.1,
        'bpr_cnt': 500,
        'icf_cnt': 500,
        'ucf_cnt': 500,
        'pop_cnt': 500,
        'methods': {'icf': 60, 'bpr': 60, 'pop': 85, 'ucf': 295}
    }

    print('Recommender: Initializing...')
    recommender = Recommender(hyper_params)
    print('Recommender: Finished.')

    start_time = time.time()
    print('Recommender: Start recommend...')
    result = recommender.recommend()
    print(f'Recommender: Finished. Time: {time.time()-start_time} seconds.')
    print(len(result[0]))

    evaluate(hyper_params['datasets'], result, cold_start=True)

    # Save result
    result_tensor = torch.empty([len(result), 500], dtype=torch.long)
    for i in range(len(result)):
        result_tensor[i] = torch.tensor(result[i])

    torch.save(result_tensor, 'prerec_result_small.pkl')
